Skip to content

fix(source-chat): emit SSE keepalives during generation and add stop-… - #1332

Open
AugustoSandim wants to merge 19 commits into
lfnovo:mainfrom
AugustoSandim:fix/source-chat-sse-keepalive-cancel
Open

fix(source-chat): emit SSE keepalives during generation and add stop-…#1332
AugustoSandim wants to merge 19 commits into
lfnovo:mainfrom
AugustoSandim:fix/source-chat-sse-keepalive-cancel

Conversation

@AugustoSandim

@AugustoSandim AugustoSandim commented Sep 5, 2026

Copy link
Copy Markdown
Contributor

Summary

The source chat SSE stream (POST /api/sources/{id}/chat/sessions/{session_id}/messages) had three gaps that together caused the reply to appear only after the user navigated away (see #1327):

  1. No keepalive — the backend yielded the user_message, then wrote nothing until the LLM finished. Slow local models mean minutes with zero bytes on the wire, so the Next.js rewrite (or any proxy) drops the idle connection.
  2. Cancel was a no-op — the frontend had an abortControllerRef but never instantiated it, so there was no way to stop an in-flight stream.
  3. No server-side stop — generation ran via asyncio.to_thread, so even after cancel the model kept burning tokens to completion.

This PR fixes all three.

What changed

Backend — keepalive + cancellable generation (api/routers/source_chat.py, open_notebook/graphs/source_chat.py)

  • Emit an SSE comment (: ping) every 15s while the model generates, so the connection never goes idle.
  • Convert the graph node and the streaming endpoint to async, and run the graph with ainvoke in a task so generation is genuinely cancellable.
  • Detect client disconnect with request.is_disconnected() and cancel the in-flight task in a finally block.
  • Persist the user message to the checkpoint up front via aupdate_state so it survives a mid-generation disconnect (the frontend refetches the checkpoint on cancel/complete).
  • Add HybridSqliteSaver, a SqliteSaver subclass with async delegates to the existing sync SQLite connection, so ainvoke/aupdate_state work without migrating the module-level sync connection.

Frontend — stop button (frontend/src/lib/api/source-chat.ts, hooks/use-source-chat.ts, components/sources/ChatPanel.tsx, page.tsx)

  • Thread an AbortSignal through sendMessage to fetch.
  • Instantiate an AbortController per message, abort the previous one when sending a new message, and abort on unmount.
  • Add a Stop button in the chat composer that cancels the in-flight stream.
  • Add the chat.stop translation to all 14 locales.

Tests

  • Characterization tests for the keepalive behavior and the disconnect-cancellation behavior.

Related Issue

Fixes #1327

Type of Change

  • Bug fix (non-breaking change that fixes an issue)

How Has This Been Tested?

  • Added new unit tests
  • Existing tests pass (uv run pytest)
  • Manual testing performed (describe below)

Test Details:

  • uv run pytest tests/ — 658 passed
  • ruff check and uv run python -m mypy — clean
  • Manual: send a message with a slow local model, confirm : ping lines stream during generation, the Stop button cancels the stream, and the reply streams/refetches when it arrives.

Design Alignment

Which design principles does this PR support? (See VISION.md)

  • Async-First for Performance

Explanation:
Generation now runs on an awaitable coroutine instead of a worker thread, so it can be cancelled the moment the client disconnects instead of running to completion.

Checklist

Code Quality

  • My code follows PEP 8 style guidelines (Python)
  • My code follows TypeScript best practices (Frontend)
  • I have added type hints to my code (Python)
  • I have performed a self-review of my code
  • I have commented my code, particularly in hard-to-understand areas
  • My changes generate no new warnings or errors

Testing

  • I have added tests that prove my fix is effective or that my feature works
  • New and existing unit tests pass locally with my changes
  • I ran linting: make ruff or ruff check . --fix
  • I ran type checking: make lint or uv run python -m mypy .

Documentation

  • I have added/updated docstrings for new/modified functions
  • I have added comments to complex logic

Additional Context

Partial text is still not streamed token-by-token — ainvoke returns the full result on completion. Streaming individual tokens would require switching to astream/astream_events, which is a separate follow-up.

Pre-Submission Verification

Before submitting, please verify:

  • I have read CONTRIBUTING.md
  • I have read VISION.md
  • This PR addresses an approved Issue assigned to me
  • I have not included unrelated changes in this PR
  • My PR title follows conventional commits format (e.g., "feat: add user authentication")

AugustoSandim and others added 3 commits September 5, 2026 15:16
…streaming button

Keep the source chat SSE connection alive while the LLM generates by sending
ignored  comments every 15 seconds, preventing proxies (including the
Next.js rewrite in front of FastAPI) from dropping the idle connection.

On the frontend, wire up an AbortController so users can stop an in-flight
stream with a new stop button, abort previous requests when sending a new
message, and abort on unmount. Add  translations for all supported
locales and a characterization test for the keepalive behavior.
… persist user message

Convert the source chat graph node and streaming endpoint to async so
generation can be cancelled when the client disconnects. Persist the
user message to the checkpoint up front via aupdate_state so it survives
a mid-generation disconnect. Add a HybridSqliteSaver to delegate
LangGraph async checkpointer calls to the existing sync SQLite connection.
Update characterization tests for the async path and add a disconnect
cancellation test.
@AugustoSandim
AugustoSandim marked this pull request as ready for review September 5, 2026 20:06

@cubic-dev-ai cubic-dev-ai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

All reported issues were addressed across 21 files

Reply with feedback, questions, or to request a fix.

Re-trigger cubic

Comment thread frontend/src/components/sources/ChatPanel.tsx Outdated
Comment thread frontend/src/lib/hooks/use-source-chat.ts Outdated
Comment thread api/routers/source_chat.py Outdated
Comment thread tests/test_chat_routers_characterization.py Outdated
… superseded streams

Check the checkpoint state before persisting the user message in
; skip the up-front append when the same human
message is already the unanswered trailing turn, preventing duplicates on
retry after a failed generation.

In the frontend, guard the streaming cleanup so a superseded send (a newer
message replaced the in-flight one) does not clear the newer stream's
loading state or trigger a stale refetch. Also show a disabled spinner in
the chat composer when streaming has no cancel callback (e.g. notebook
chat) instead of a non-functional Stop button.

Update characterization tests: make the keepalive test deterministic with
an , and add tests for the duplicate-pending-message skip
and normal append-after-AI-message behavior.

@cubic-dev-ai cubic-dev-ai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

All reported issues were addressed across 4 files (changes from recent commits).

Reply with feedback, questions, or to request a fix.

Re-trigger cubic

Comment thread api/routers/source_chat.py Outdated
Comment thread frontend/src/lib/hooks/use-source-chat.ts
Comment thread api/routers/source_chat.py Outdated
Comment thread frontend/src/lib/hooks/use-source-chat.ts Outdated
Comment thread api/routers/source_chat.py Outdated
…-session streams

Add a client-generated  to  and use it (instead of content equality) to skip the up-front user-message append on retry, while still keeping distinct identical messages. Serialize the snapshot → append → invoke sequence per session with an  so concurrent requests for the same thread cannot start duplicate generations.

Update the frontend to reuse the trailing unanswered human message id when retrying the same content and generate a fresh uuid otherwise, and to clear streaming state when session creation fails. Add ADR-009 documenting the  async-to-sync bridge. Add/update characterization and hook tests for message-id dedup, distinct identical messages, and streaming lifecycle.
@AugustoSandim
AugustoSandim marked this pull request as draft September 5, 2026 21:19

@cubic-dev-ai cubic-dev-ai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

1 issue found across 7 files (changes from recent commits).

Confidence score: 4/5

  • frontend/src/lib/hooks/use-source-chat.ts couples retry/deduplication and message-ID selection to React state, making security-sensitive behavior harder to test and maintain; extract the selection policy into a pure, directly testable function.
Prompt for AI agents (unresolved issues)

Check if these issues are valid — if so, understand the root cause of each and fix them. If appropriate, use sub-agents to investigate and fix each issue separately.


<file name="frontend/src/lib/hooks/use-source-chat.ts">

<violation number="1" location="frontend/src/lib/hooks/use-source-chat.ts:152">
P2: Custom agent: **Security & testability**

The retry/deduplication policy now lives inside `useSourceChat`, where it is coupled to React state instead of being directly testable. Move message-ID selection into a pure service/domain utility and test both retry reuse and distinct identical messages there.</violation>
</file>

Requires human review: Auto-approval blocked because this review re-detected 1 unresolved issue already reported by Cubic.
Tip: Review your code locally with the cubic CLI to iterate faster.

Re-trigger cubic

Comment thread frontend/src/lib/hooks/use-source-chat.ts Outdated
Comment thread api/routers/source_chat.py Outdated
Comment thread api/routers/source_chat.py Outdated
Comment thread docs/7-DEVELOPMENT/decisions/ADR-009-hybrid-sqlite-checkpointer-bridge.md Outdated
Comment thread frontend/src/lib/hooks/use-source-chat.ts Outdated
The test fixture set `model_override: null`, violating
`SourceChatSession.model_override?: string` (the interface narrows the
base type's `string | null`). And `sendPromise` was read before
TypeScript could prove it assigned (assigned only inside an `act()`
callback), triggering TS2454. Drop the invalid field and use a
definite-assignment assertion.

Co-authored-by: Cursor <cursoragent@cursor.com>
@AugustoSandim
AugustoSandim force-pushed the fix/source-chat-sse-keepalive-cancel branch from 05121ca to 852c1c9 Compare September 5, 2026 22:00
…sage ids

- Replace the global per-session `asyncio.Lock` dict with a refcounted `_SessionLock` that evicts its entry once the last stream releases, preventing an unbounded lock table in a long-lived process.
- Add characterization test verifying the lock entry is removed after the stream finishes.
- Extract `selectMessageId()` to decide whether to reuse the trailing unanswered human id (retry) or generate a fresh uuid (new turn), and send that real id optimistically so the backend `already_pending` check matches the optimistic entry.
- Stop filtering `temp-*` ids on send error since optimistic messages now carry their real ids.
- Update ADR-009 to include `aget_tuple` in the async-to-sync delegation surface.
@AugustoSandim
AugustoSandim marked this pull request as ready for review September 5, 2026 22:19

@cubic-dev-ai cubic-dev-ai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

All reported issues were addressed across 27 files

Requires human review: Auto-approval blocked because this review re-detected 1 unresolved issue already reported by Cubic.

Re-trigger cubic

Comment thread frontend/src/lib/hooks/use-source-chat.ts
Comment thread frontend/src/lib/hooks/use-source-chat.ts
Comment thread frontend/src/lib/hooks/use-source-chat.ts
Comment thread api/routers/source_chat.py Outdated
Comment thread api/routers/source_chat.py Outdated
Comment thread frontend/src/components/sources/ChatPanel.tsx Outdated
…d lock waits

- Add a 1s disconnect poll interval separate from the 15s SSE keepalive so dropped connections stop generation promptly; only emit  keepalives when due.
- Track whether the session lock was actually acquired and, if the stream is cancelled while waiting, decrement the holder count without releasing an unheld lock.
- Serialize concurrent first-time sends into a single session-create promise, avoid duplicate optimistic bubbles on retry, and roll back the optimistic user message if send fails before persistence.
- Use a  label for the non-cancellable notebook-chat spinner and add translations for all locales.
- Update characterization and hook tests for the new poll interval, lock cancellation, and send races.

@cubic-dev-ai cubic-dev-ai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

All reported issues were addressed across 23 files (changes from recent commits).

Tip: Review your code locally with the cubic CLI to iterate faster.

Re-trigger cubic

Comment thread frontend/src/lib/hooks/use-source-chat.ts
Comment thread frontend/src/lib/utils/source-chat-message.ts Outdated
Comment thread frontend/src/lib/hooks/use-source-chat.ts Outdated
Comment thread frontend/src/components/sources/ChatPanel.test.tsx Outdated
@AugustoSandim
AugustoSandim marked this pull request as draft September 6, 2026 18:22
…bles on failure, and support non-secure UUID generation

- When a user cancels streaming while an auto-create session request is still in flight, adopt the created session id once it resolves and invalidate the sessions list so the next send does not create another empty session.
- On send failure, only roll back an optimistic user message added in the current turn; a retry that reuses a persisted trailing human id keeps its bubble visible until the refetch completes.
- Add `createMessageId()` with a `Math.random` fallback for contexts where `crypto.randomUUID` is unavailable (e.g. non-secure HTTP) and use it as the default id generator.
- Update tests for the new stop-during-create behavior, retry-failure visibility, and UUID fallback; rename the notebook-chat composer test to clarify "without stop support".
@AugustoSandim
AugustoSandim marked this pull request as ready for review September 6, 2026 18:48

@cubic-dev-ai cubic-dev-ai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

All reported issues were addressed across 28 files

Reply with feedback, questions, or to request a fix.

Re-trigger cubic

Comment thread frontend/src/lib/utils/source-chat-message.ts
Comment thread frontend/src/lib/hooks/use-source-chat.ts Outdated
Comment thread frontend/src/lib/hooks/use-source-chat.ts
Comment thread api/routers/source_chat.py Outdated
…tend stale-state races

- Move the per-session lock and pending human-turn persistence policy from `api/routers/source_chat.py` into a new `api/source_chat_service.py` module, exposing `source_chat_turn` so the router delegates serialization and persistence.
- Add focused unit tests for the refcounted session lock (serialization, eviction, and cancelled-waiter cleanup) and the message-id-based pending-turn deduplication; update characterization tests to import from the new module.
- In the frontend `useSourceChat` hook, use refs and a per-send generation token so stale session snapshots cannot overwrite messages from a newer stream, resolve authoritative state before choosing the message id, and skip session adoption when the abort comes from unmount rather than Stop.
- Add tests for authoritative state resolution, stale-refetch suppression, and unmount abort behavior.

@cubic-dev-ai cubic-dev-ai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

All reported issues were addressed across 9 files (changes from recent commits).

Requires human review: Auto-approval blocked because this review re-detected 1 unresolved issue already reported by Cubic.
Tip: Review your code locally with the cubic CLI to iterate faster.

Re-trigger cubic

Comment thread frontend/src/lib/hooks/use-source-chat.ts Outdated
Comment thread frontend/src/lib/hooks/use-source-chat.ts Outdated
…ssages across session switches

- Add `getLogSafeErrorMessage()` to strip `Authorization` headers and bound error length before `console.error`, so axios request configs carrying bearer tokens are never logged.
- Use the new helper in `useSourceChat` for session creation, hydration, and send failures.
- Track which session the shared `messages` list represents via `messagesSessionRef`; only apply stream chunks and optimistic turns when the list still belongs to the originating session, and re-attach the full accumulated answer when the user switches back.
- Add tests verifying auth tokens are absent from logs, sends stay on their original session during pre-send hydration switches, and streaming sessions rehydrate correctly after switching away and back.

@cubic-dev-ai cubic-dev-ai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

All reported issues were addressed across 4 files (changes from recent commits).

Tip: Review your code locally with the cubic CLI to iterate faster.

Re-trigger cubic

Comment thread frontend/src/lib/hooks/use-source-chat.ts Outdated
Comment thread frontend/src/lib/hooks/use-source-chat.test.tsx Outdated
@AugustoSandim
AugustoSandim marked this pull request as draft September 6, 2026 20:13
  before claiming the shared list and restore console spy
  safely
@AugustoSandim
AugustoSandim marked this pull request as ready for review September 6, 2026 20:40

@cubic-dev-ai cubic-dev-ai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

All reported issues were addressed across 32 files

Tip: cubic can generate docs of your entire codebase and keep them up to date. Try it here.

Re-trigger cubic

Comment thread frontend/src/lib/hooks/use-source-chat.ts
Comment thread api/routers/source_chat.py Outdated
Comment thread api/routers/source_chat.py Outdated
Comment thread frontend/src/lib/hooks/use-source-chat.ts
Comment thread frontend/src/lib/hooks/use-source-chat.ts Outdated
Comment thread frontend/src/lib/hooks/use-source-chat.test.tsx Outdated
…superseded stream handling

- Add `SessionTurnLease` in `api/source_chat_service.py` to acquire the per-session turn lock by polling instead of blocking, so SSE streams can emit keepalive comments and detect client disconnects while queued behind another turn.
- Update `stream_source_chat_response` to poll the lease, emit due keepalives before the final response, and release the lease cleanly on cancellation.
- Fail sends in `useSourceChat` when pre-send session hydration errors, rather than minting an unverified message id that would bypass backend dedup.
- Ensure only the latest send writes shared state: guard `context_indicators` and message ownership against superseded streams, and refetch the persisted checkpoint after cancellation so the pending turn survives.
- Add characterization tests for queued keepalives, queued disconnect cleanup, and due keepalive before completion; update hook tests for cancellation unwind, retry id reuse, and hydration failure.

@cubic-dev-ai cubic-dev-ai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

3 issues found across 5 files (changes from recent commits).

Confidence score: 3/5

  • frontend/src/lib/hooks/use-source-chat.ts does not claim the newly created session’s empty message list, so ownsMessages() can reject the optimistic bubble and streamed answer while the session query loads; update this branch to claim the new session’s list.
  • frontend/src/lib/hooks/use-source-chat.test.tsx restores its console.error spy only around the later act and assertions, so an earlier renderHook or waitFor failure can leak the spy into other tests; wrap the full setup and wait sequence in guaranteed cleanup.
  • tests/test_chat_routers_characterization.py assumes exactly two time.monotonic() reads, making the keepalive test brittle if control flow changes; use a clock stub or sequence that tolerates the expected additional reads.
Prompt for AI agents (unresolved issues)

Check if these issues are valid — if so, understand the root cause of each and fix them. If appropriate, use sub-agents to investigate and fix each issue separately.


<file name="frontend/src/lib/hooks/use-source-chat.ts">

<violation number="1" location="frontend/src/lib/hooks/use-source-chat.ts:297">
P2: When the first message auto-creates a session, this branch never claims the new session’s empty list, so `ownsMessages()` rejects the optimistic bubble and streamed answer while the session query is still loading. Adopt an empty list for `sessionJustCreated` before sending so the first turn is visible immediately and does not depend on query timing.</violation>
</file>

<file name="frontend/src/lib/hooks/use-source-chat.test.tsx">

<violation number="1" location="frontend/src/lib/hooks/use-source-chat.test.tsx:240">
P3: The `console.error` spy is installed before `renderHook`/`waitFor`, but the `try/finally` that restores it only wraps the later `act` and assertions. If the `waitFor(() => expect(result.current.currentSessionId).toBe('session:1'))` assertion fails, the test throws before entering the `try`, the spy is never restored, and `console.error` stays silenced for every test that follows — the exact leak the sibling 'never logs the request config' test was changed to avoid in this same PR. Move the `renderHook`, `waitFor`, and `act` inside the `try` block so the `finally` always restores the spy.</violation>
</file>

<file name="tests/test_chat_routers_characterization.py">

<violation number="1" location="tests/test_chat_routers_characterization.py:716">
P3: The fake clock in `test_stream_source_chat_emits_due_keepalive_before_final_response` hard-codes exactly two `time.monotonic()` reads (`clock = iter([0.0, 100.0])`). It works only because the current control flow reads the clock once for `last_keepalive` initialization and once in the generation loop, and because the session lock is free so the queue phase never runs. Any legitimate change to the streaming flow that adds a third clock read (an extra keepalive check, a second poll iteration, or a disconnect probe) exhausts the iterator and fails the test with an opaque `StopIteration` before any assertion runs. Give the fake clock a non-raising implementation so future control-flow changes produce a meaningful assertion failure instead of a confusing exception.</violation>
</file>

Tip: Review your code locally with the cubic CLI to iterate faster.

Re-trigger cubic

Comment thread frontend/src/lib/hooks/use-source-chat.ts Outdated
Comment thread frontend/src/lib/hooks/use-source-chat.test.tsx Outdated
Comment thread tests/test_chat_routers_characterization.py Outdated
AugustoSandim and others added 4 commits September 6, 2026 18:21
Co-authored-by: cubic-dev-ai[bot] <191113872+cubic-dev-ai[bot]@users.noreply.github.com>
Co-authored-by: cubic-dev-ai[bot] <191113872+cubic-dev-ai[bot]@users.noreply.github.com>
Co-authored-by: cubic-dev-ai[bot] <191113872+cubic-dev-ai[bot]@users.noreply.github.com>
…ot edit

A cubic-dev-ai suggested edit duplicated the try/finally in the
"fails the send when pre-send hydration errors" test, leaving a dangling
second block that referenced `result` out of scope. It broke tsc, the
frontend build, and the test run (ReferenceError: result is not defined).

Co-Authored-By: Claude <noreply@anthropic.com>

@cubic-dev-ai cubic-dev-ai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

All reported issues were addressed across 3 files (changes from recent commits).

Tip: Review your code locally with the cubic CLI to iterate faster.

Re-trigger cubic

Comment thread tests/test_chat_routers_characterization.py Outdated
Comment thread frontend/src/lib/hooks/use-source-chat.ts
AugustoSandim and others added 2 commits September 6, 2026 18:52
…nd hydration

The session-created guard left `else if (!sessionJustCreated)` dead after
the explicit fresh-session branch was added; fold it into a plain `else`.

Co-Authored-By: Claude <noreply@anthropic.com>
…ive test

The previous counter froze the clock at 100.0 for every call after the
first, which left the due-keepalive check at 0 forever if the generation
loop ever iterated twice. A monotonically increasing clock keeps the
keepalive-before-final-response regression test meaningful under any
number of polls.

Co-Authored-By: Claude <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Source chat SSE has no keepalive and cancel is a no-op

1 participant